#include "../src/mstep.hh"
#include "../src/distribution.cc"
#include "../src/nominal.hh"

#include <cassert>
#include <iostream>

namespace
{
using namespace playerpiano;

struct MstepTestDatum 
{ 
  struct
    {
      float    alpha[4 * 5 * 5];
      float    logbeta;
      float    pz[5];
      Rating   ratings[8];
      float    logbetamean;
      float    logbetastddev;
    }                   inputs;
  struct
    {
      float    q;
      float    logbeta;
    }                   desired_output;
};

static void
test_mstep (void)
{
  MstepTestDatum data[] = 
    {{{{-1., 0.483197, 0.660258, 0.727927, 0.982813, 0.767098, -1.,  0.675677, 0.447308, 0.616571, 0.197368, 0.233449, -1., 0.0767365,  0.180397, 0.886653, 0.303462, 0.853033, -1., 0.516409, 0.143031,  0.0708752, 0.577288, 0.982266, -1., -1., 0.592662, 0.756706,  0.0143276, 0.802102, 0.109465, -1., 0.0964475, 0.2864, 0.819289,  0.342367, 0.42077, -1., 0.839092, 0.202718, 0.144999, 0.187321,  0.762356, -1., 0.0223215, 0.258346, 0.883859, 0.909323,  0.505913, -1., -1., 0.115314, 0.812984, 0.332035, 0.523646,  0.522653, -1., 0.0562783, 0.317707, 0.721544, 0.413188,  0.959831, -1., 0.0313071, 0.902255, 0.0708207, 0.53906,  0.192215, -1., 0.699537, 0.925822, 0.351739, 0.429859,  0.677215, -1., -1., 0.667476, 0.46788, 0.520537, 0.171302,  0.552161, -1., 0.654896, 0.188502, 0.647656, 0.0295083,  0.598618, -1., 0.870794, 0.926112, 0.61632, 0.638787,  0.839487, -1., 0.0238569, 0.5455, 0.0997268, 0.647272,  0.32432, -1.},  0., {0.619678, 0.747988, 0.217413, 0.647105,  0.952202}, {{0, 4}, {0, 0}, {1, 0}, {2, 2}, {3, 2}, {2, 1}, {0,  1}, {1, 2}}, 0.280108,  0.696876}, {-42.867210538589882660, -0.84826212415227818426}},  {{{-1., 0.475803, 0.400041, 0.625211, 0.508375, 0.828147, -1.,  0.370533, 0.0265936, 0.63758, 0.902035, 0.754213, -1., 0.387806,  0.798093, 0.878178, 0.208713, 0.28808, -1., 0.150821, 0.553858,  0.589035, 0.540092, 0.933407, -1., -1., 0.906753, 0.636833,  0.259984, 0.236531, 0.43095, -1., 0.236792, 0.634773, 0.728156,  0.602803, 0.866259, -1., 0.608179, 0.0905762, 0.700769, 0.112046,  0.220373, -1., 0.292483, 0.822591, 0.903333, 0.932293,  0.141663, -1., -1., 0.268733, 0.314298, 0.392201, 0.208255,  0.36198, -1., 0.677465, 0.132217, 0.971724, 0.93103,  0.440674, -1., 0.497444, 0.243568, 0.328226, 0.574415,  0.889265, -1., 0.152992, 0.627458, 0.462369, 0.668892,  0.860508, -1., -1., 0.804867, 0.559035, 0.736598, 0.718846,  0.536134, -1., 0.244737, 0.344397, 0.51059, 0.174154,  0.567271, -1., 0.21218, 0.538866, 0.243124, 0.126598,  0.714737, -1., 0.295298, 0.914898, 0.552183, 0.825472,  0.142307, -1.},  0., {0.28744, 0.0898143, 0.15658, 0.281798,  0.482573}, {{2, 3}, {2, 4}, {0, 2}, {1, 2}, {0, 0}, {3, 4}, {0,  4}, {0, 3}}, 0.530779,  0.419982}, {-17.771252954779590711, -0.11887234805451569553}},  {{{-1., 0.562952, 0.946439, 0.286042, 0.0755847, 0.0523619, -1.,  0.772285, 0.718771, 0.863404, 0.513496, 0.52916, -1., 0.592173,  0.148668, 0.218197, 0.614262, 0.0399903, -1., 0.323196, 0.0758907,  0.326822, 0.950176, 0.166615, -1., -1., 0.794093, 0.844249,  0.419397, 0.746633, 0.23114, -1., 0.89781, 0.133355, 0.671049,  0.178778, 0.125525, -1., 0.414584, 0.807644, 0.665283, 0.596364,  0.822411, -1., 0.658977, 0.447085, 0.982102, 0.78242,  0.335781, -1., -1., 0.371195, 0.65528, 0.832244, 0.169166,  0.577102, -1., 0.811031, 0.412847, 0.422533, 0.345962,  0.913222, -1., 0.279493, 0.751484, 0.167184, 0.787697,  0.864909, -1., 0.94384, 0.501901, 0.191333, 0.042498,  0.284863, -1., -1., 0.0548158, 0.209231, 0.260078, 0.949082,  0.683621, -1., 0.553951, 0.427833, 0.779917, 0.106519,  0.74292, -1., 0.0149857, 0.357384, 0.760557, 0.829698,  0.735493, -1., 0.6059, 0.593373, 0.042001, 0.870584,  0.66206, -1.},  0., {0.091472, 0.850668, 0.828086, 0.377197,  0.0366562}, {{2, 4}, {3, 2}, {0, 3}, {1, 3}, {0, 4}, {3, 3}, {2,  4}, {2, 1}}, 0.641438,  0.568009}, {-31.565851431750928371, -0.52261759603711178228}},  {{{-1., 0.428115, 0.353035, 0.0874867, 0.140175, 0.648198, -1.,  0.246516, 0.344567, 0.12519, 0.290814, 0.485959, -1., 0.514869,  0.389697, 0.684914, 0.892586, 0.472868, -1., 0.519113, 0.022854,  0.801114, 0.6222, 0.691026, -1., -1., 0.645657, 0.764458,  0.980762, 0.123018, 0.217542, -1., 0.411423, 0.893275, 0.982842,  0.569344, 0.164907, -1., 0.548708, 0.857652, 0.27853, 0.678947,  0.0338393, -1., 0.467956, 0.593615, 0.786361, 0.560971,  0.948843, -1., -1., 0.570761, 0.985247, 0.938772, 0.257817,  0.925104, -1., 0.220789, 0.958009, 0.134799, 0.707562,  0.809366, -1., 0.0647339, 0.151957, 0.138218, 0.64446,  0.516025, -1., 0.294305, 0.859689, 0.965512, 0.482186,  0.826349, -1., -1., 0.266074, 0.179151, 0.921215, 0.877506,  0.695312, -1., 0.193904, 0.982443, 0.619689, 0.770208,  0.973115, -1., 0.024434, 0.484889, 0.0626459, 0.163749,  0.9597, -1., 0.332932, 0.924427, 0.519289, 0.443675,  0.0386276, -1.},  0., {0.0647385, 0.553777, 0.961489, 0.212279,  0.798665}, {{2, 1}, {3, 1}, {3, 0}, {3, 4}, {1, 3}, {0, 3}, {2,  2}, {1, 0}}, 0.374626, 0.0402736}, {-44.998991806421287124,  0.33303184067232409168}}}; 

  for (unsigned int i = 0; i < sizeof (data) / sizeof (data[0]); ++i)
    {
      NormalDistribution prior_logbeta (data[i].inputs.logbetamean, 
                                        data[i].inputs.logbetastddev);
      Nominal nominal (data[i].inputs.alpha, 0, 0, 4, 5, 0);

      for (unsigned int j = 0; j < 4; ++j)
        {
          for (unsigned int k = 0; k < 5; ++k)
            {
              data[i].inputs.alpha[j * 5 * 5 + k * 5 + k] = drand48 ();
            }
        }

      float hq = 0;

      for (unsigned int k = 0; k < 5; ++k)
        {
          hq -= data[i].inputs.pz[k] * log (data[i].inputs.pz[k]);
        }

      std::vector<float> pz (5);
      std::vector<Rating> ratings (8);

      for (unsigned int k = 0; k < 5; ++k)
        {
          pz[k] = data[i].inputs.pz[k];
        }

      for (unsigned int k = 0; k < 8; ++k)
        {
          ratings[k] = data[i].inputs.ratings[k];
        }

      float logbeta = data[i].inputs.logbeta;

      float q = mstep (nominal,
                       logbeta,
                       pz,
                       ratings,
                       prior_logbeta,
                       hq,
                       1000,
                       1e-7);

      assert (fabs (q - data[i].desired_output.q) <=
              1e-4 * (1 + fabs (q) + fabs (data[i].desired_output.q)) ||
              (std::cerr << q << " ?= " << data[i].desired_output.q << std::endl,
               0));

      assert (fabs (logbeta - data[i].desired_output.logbeta) <=
              5e-3 * (1 + fabs (logbeta) + fabs (data[i].desired_output.logbeta)) ||
              (std::cerr << logbeta << " ?= " << data[i].desired_output.logbeta << std::endl,
               0));

    }
}

}

int 
main (void) 
{
  srand48 (69);

  test_mstep ();

  return 0;
}
